from ultralytics import YOLO
# import torch

if __name__ == "__main__":
    # Load the model.
    model = YOLO('best.pt')
    model.to('cuda')



    results = model.train(data='D:/DIPLOM/YOLODataset/dataset.yaml', imgsz=640, epochs=300, batch=6, name='yolov8nx_custom')

    #model.tune(data='D:/DIPLOM/YOLODataset/dataset.yaml', epochs=5, iterations=5, optimizer='AdamW', plots=False, save=False, val=False)